import numpy as np

class ExtractionSolver:
    def __init__(self):
        self.initialize_params()

    def initialize_params(self):
        """初始化计算参数"""
        self.stage = 5
        self.stream = int(self.stage) + 1
        self.R = np.zeros(self.stream)
        self.E = np.zeros(self.stream)
        self.r = np.zeros(self.stream)
        self.e = np.zeros(self.stream)
        self.K = np.zeros(self.stream)

        # 初始条件
        self.r_out = 0.05
        self.e_in = 0
        self.r[0] = 0.5
        self.e[0] = self.e_in
        self.R[:] = 100
        self.E[:] = 50
        self.K[:] = 1.5
        self.k = 1.5
        self.sp = 1.5

    def cross(self, R_opt):
        """错流萃取计算函数"""
        n = self.stream
        # 跳过第一个元素更新
        self.R[1:] = R_opt
        # 流股物料平衡
        for j in range(0, 200):
            for i in range(1, n):
                self.E[i] = self.E[0] + self.R[i-1] * self.r[i-1] - self.R[i] * self.r[i]
                self.R[i] = self.R[i-1] - self.R[i-1] * self.r[i-1] + self.R[i] * self.r[i]
            self.r[1:] = (self.R[:n-1] * self.r[:n-1] + self.E[0] * self.e[0]) / (self.E[1:] * self.K[1:] + self.R[1:])
            # 更新分配系数
            self.K = self.k * self.r + self.sp
            # 更新萃余相浓度
            self.e[1:] = self.K[1:] * self.r[1:]
            # 关键物料的平衡方程
            f = self.R[1:] - self.R[:n-1] + self.R[:n-1] * self.r[:n-1] - self.R[1:] * self.r[1:]
            if sum(f ** 2) < 1e-5:
                break
            if j > 100:
                return False
        return True

    def counter(self, R_opt):
        """逆流萃取计算函数"""
        n = self.stream
        # 跳过第一个元素更新
        self.R[1:] = R_opt
        # 更新萃余相的流量
        for j in range(0, 200):
            for i in range(1, n):
                self.R[i] = self.R[i-1] - (self.R[i-1] * self.r[i-1] - self.R[i] * self.r[i])
            # 更新萃取相的流量
            self.E[n-1] = self.E[0] + self.R[n-2] * self.r[n-2] - self.R[n-1] * self.r[n-1]
            for i in range(n-2, 0, -1):
                self.E[i] = self.E[i+1] + self.R[i-1] * self.r[i-1] - self.R[i] * self.r[i]
            # 更新分配系数
            self.K = self.k * self.r + self.sp
            # 更新萃取相的浓度
            self.e[1:] = self.K[1:] * self.r[1:]
            # 更新萃余相的浓度
            self.r[n-1] = (self.R[n-2] * self.r[n-2] + self.E[0] * self.e[0]) / (self.E[n-1] * self.K[n-1] + self.R[n-1])
            for i in range(1, n-1):
                self.r[i] = (self.R[i-1] * self.r[i-1] + self.E[i+1] * self.e[i+1]) / (self.E[i] * self.K[i] + self.R[i])
            # 关键物料的平衡方程
            f = self.R[:n-1] - self.R[1:] - (self.R[:n-1] * self.r[:n-1] - self.R[1:] * self.r[1:])
            if sum(f ** 2) < 1e-5:
                break
            if j > 100:
                return False
        return True

    # 求错流萃取/逆流萃取
    def solve_cross(self):
        """求解萃取理论级"""
        try:
            return self.cross(self.R[1:])
        except Exception as e:
            print(f"计算出错: {e}")
            return False

    def solve_counter(self):
        """求解逆流理论级"""
        try:
            return self.counter(self.R[1:])
        except Exception as e:
            print(f"计算出错: {e}")
            return False

    def show_results(self, is_counter):
        """生成详细的结果表格"""
        result = []
        result.append(f"{'级数':<5}{'分配系数':<8}{'萃余相浓度':<10}{'萃取相浓度':<10}{'萃余相流量':<12}{'萃取相流量':<15}\n")
        for i in range(0, self.stream):
            result.append(f"{i:<7}{self.K[i]:<12.4f}{self.r[i]:<15.4f}{self.e[i]:<15.4f}{self.R[i]:<17.4f}{self.E[i]:<10.4f}\n")
        for i in range(0, self.stream):
            if self.r[i] < self.r_out:
                if not is_counter:
                    result.append(f"需要的萃取级数为{i}级！")
                else:
                    result.append(f"可以满足要求！")
                break
            elif i == self.stream - 1:
                result.append(f"无法达到萃取要求！")
        print(''.join(result))
        return True

solve = ExtractionSolver()
result = solve.solve_cross()
solve.show_results(0)